import json
import paramiko
import threading
from channels.generic.websocket import WebsocketConsumer
from .models import Host


class SSHConsumer(WebsocketConsumer):
    """SSH终端WebSocket消费者"""
    
    def connect(self):
        """建立WebSocket连接"""
        self.accept()
        self.ssh_client = None
        self.ssh_channel = None
        
    def disconnect(self, close_code):
        """断开连接"""
        if self.ssh_channel:
            self.ssh_channel.close()
        if self.ssh_client:
            self.ssh_client.close()
    
    def receive(self, text_data):
        """接收消息"""
        try:
            data = json.loads(text_data)
            action = data.get('action')
            
            if action == 'connect':
                # 建立SSH连接
                host_id = data.get('host_id')
                self.connect_ssh(host_id)
            elif action == 'input':
                # 发送输入到SSH
                command = data.get('data', '')
                if self.ssh_channel:
                    self.ssh_channel.send(command)
            elif action == 'resize':
                # 调整终端大小
                cols = data.get('cols', 80)
                rows = data.get('rows', 24)
                if self.ssh_channel:
                    self.ssh_channel.resize_pty(width=cols, height=rows)
        except Exception as e:
            self.send(text_data=json.dumps({
                'error': str(e)
            }))
    
    def connect_ssh(self, host_id):
        """连接SSH"""
        try:
            host = Host.objects.get(id=host_id)
            
            # 创建SSH客户端
            self.ssh_client = paramiko.SSHClient()
            self.ssh_client.set_missing_host_key_policy(paramiko.AutoAddPolicy())
            
            # 连接
            if host.private_key:
                from io import StringIO
                key = paramiko.RSAKey.from_private_key(StringIO(host.private_key))
                self.ssh_client.connect(
                    hostname=host.ip_address,
                    port=host.port,
                    username=host.username,
                    pkey=key,
                    timeout=10
                )
            else:
                self.ssh_client.connect(
                    hostname=host.ip_address,
                    port=host.port,
                    username=host.username,
                    password=host.password,
                    timeout=10
                )
            
            # 创建交互式shell
            self.ssh_channel = self.ssh_client.invoke_shell(term='xterm')
            self.ssh_channel.settimeout(0.0)
            
            # 发送连接成功消息
            self.send(text_data=json.dumps({
                'status': 'connected',
                'message': f'已连接到 {host.hostname} ({host.ip_address})'
            }))
            
            # 启动线程读取SSH输出
            threading.Thread(target=self.read_ssh_output, daemon=True).start()
            
        except Host.DoesNotExist:
            self.send(text_data=json.dumps({
                'error': '主机不存在'
            }))
        except Exception as e:
            self.send(text_data=json.dumps({
                'error': f'连接失败: {str(e)}'
            }))
    
    def read_ssh_output(self):
        """读取SSH输出"""
        while True:
            try:
                if self.ssh_channel and self.ssh_channel.recv_ready():
                    data = self.ssh_channel.recv(1024)
                    if data:
                        self.send(text_data=json.dumps({
                            'output': data.decode('utf-8', errors='ignore')
                        }))
                    else:
                        break
            except Exception as e:
                self.send(text_data=json.dumps({
                    'error': str(e)
                }))
                break
